from motor.motor_asyncio import AsyncIOMotorClient
from core.config import settings

class MongoDB:
    def __init__(self):
        self.client = None
        self.db = None

    # 在 MongoDB 类的 connect 方法中添加连接池配置
    async def connect(self):
        """
        连接MongoDB数据库
        """
        self.client = AsyncIOMotorClient(
            settings.MONGO_URI,
            maxPoolSize=50,  # 最大连接数
            minPoolSize=10,  # 最小连接数
            maxIdleTimeMS=30000,  # 连接最大空闲时间
            socketTimeoutMS=20000,  # 套接字超时
            connectTimeoutMS=20000,  # 连接超时
            serverSelectionTimeoutMS=20000,  # 服务器选择超时
            waitQueueTimeoutMS=10000  # 等待队列超时
        )
        self.db = self.client[settings.MONGO_DB_NAME]

    async def close(self):
        """
        关闭MongoDB连接
        """
        if self.client:
            self.client.close()

    async def check_connection(self):
        """
        检查MongoDB连接状态
        """
        try:
            await self.db.command('ping')
            return True
        except Exception as e:
            print(f"MongoDB connection check failed: {str(e)}")
            return False

    async def save_chat_message(self, message_data: dict):
        """保存聊天记录"""
        collection = self.db.chat_messages
        result = await collection.insert_one(message_data)
        return result.inserted_id

    async def get_recent_chat_history(self, user_id: str, session_id: str, limit: int = 5):
        """获取最近的聊天记录"""
        collection = self.db.chat_messages
        cursor = collection.find({
            "userId": user_id,
            "sessionId": session_id
        }).sort("createdAt", -1).limit(limit)
        
        messages = []
        async for doc in cursor:
            messages.append({
                "role": doc["role"],
                "content": doc["content"]
            })
        
        return list(reversed(messages))  # 保持时间顺序
    # 根据userId和sessionId获取聊天记录
    async def get_chat_history(self, user_id: str, session_id: str):
        """获取聊天记录"""
        collection = self.db.chat_messages
        cursor = collection.find({
            "userId": user_id,
            "sessionId": session_id
        }).sort("createdAt", 1)  
    
        messages = []
        async for doc in cursor:
            message = {
                "role": doc["role"],
                "content": doc["content"],
                "createdAt": doc["createdAt"]
            }
            # 安全地获取username字段，如果不存在则设为空字符串
            if "username" in doc:
                message["username"] = doc["username"]
            else:
                message["username"] = ""
                
            messages.append(message)

        return messages
    # 根据userId和sessionId清空聊天记录
    async def clear_chat_history(self, user_id: str, session_id: str):
        """清空聊天记录"""
        collection = self.db.chat_messages
        result = await collection.delete_many({
            "userId": user_id,
            "sessionId": session_id
        })
        return result.deleted_count

    async def get_user_all_sessions(self, user_id: str):
        """获取用户的所有对话记录，按会话分组"""
        collection = self.db.chat_messages
        pipeline = [
            {"$match": {"userId": user_id}},
            {"$sort": {"createdAt": -1}},  # 按时间倒序排列
            {"$group": {
                "_id": "$sessionId",
                "sessionId": {"$first": "$sessionId"},
                "username": {"$first": "$username"},
                "lastMessage": {"$first": "$content"},
                "lastMessageTime": {"$first": "$createdAt"},
                "messageCount": {"$sum": 1}
            }},
            {"$sort": {"lastMessageTime": -1}}  # 按最后消息时间倒序排列
        ]
        
        cursor = collection.aggregate(pipeline)
        sessions = []
        async for doc in cursor:
            sessions.append({
                "sessionId": doc["sessionId"],
                "username": doc["username"],
                "lastMessage": doc["lastMessage"],
                "lastMessageTime": doc["lastMessageTime"],
                "messageCount": doc["messageCount"]
            })
        
        return sessions

# 创建全局MongoDB实例
mongo = MongoDB()